package kodkod.analysis.AST;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.TreeSet;

import kodkod.ast.BinaryExpression;
import kodkod.ast.Node;
import kodkod.ast.UnaryExpression;
import kodkod.ast.operator.ExprOperator;

public abstract class Graph {

	HashMap<wNode, wNode> allNodes = new HashMap<wNode, wNode>();
	TreeSet<wEdge> allEdges = new TreeSet<wEdge>();
	
	HashMap<wEdge, wEdge> edgeFinder = new HashMap<wEdge, wEdge>();
	
	public HashMap<Node, Integer> nodeIds = new HashMap<Node, Integer>();
	public int idCounter = 0;
	
	public TreeSet<wNode> primaryNodes = new TreeSet<wNode>();
	public TreeSet<wNode> leafNodes = new TreeSet<wNode>();
	
	ArrayList<TreeSet<wNode>> nodeLists = new ArrayList<TreeSet<wNode>>();
	
	public Graph() {
		nodeLists.add(primaryNodes);
		nodeLists.add(leafNodes);
	}
	
	public wNode getNode(Node node) {
		wNode wn = new wNode(this, node);
		return getNode(wn);
	}
	
	public wNode getNode(wNode wn) {
		wNode node = allNodes.get(wn);
		if (node == null) {
			allNodes.put(wn, wn);
			return wn;
		} else {
			return node;
		}
	}
	
	public void replaceNode(wNode oldNode, wNode newNode) {
		allNodes.put(oldNode, newNode);
		
		// patch node-lists
		for (TreeSet<wNode> nodeList : nodeLists) {
			if (nodeList.remove(oldNode)) {
				nodeList.add(newNode);
			}
		}
		
		// remove edges
		TreeSet<wEdge> edgesToRemove = new TreeSet<wEdge>();
		edgesToRemove.addAll(oldNode.getIncomingEdges());
		edgesToRemove.addAll(oldNode.getOutgoingEdges());
		removeEdges(edgesToRemove);
	}
	
	public wEdge newEdge(wNode node1, wNode node2) {
		wEdge edge = new wEdge(node1, node2);
		allEdges.add(edge);
		edgeFinder.put(edge, edge);
		return edge;
	}
	
	public wEdge findEdge(wNode node1, wNode node2) {
		wEdge edge = new wEdge(node1, node2);
		wEdge existing = edgeFinder.get(edge);
		if (existing == null) {
			return newEdge(node1, node2);
		} else {
			return existing;
		}
	}
	
	public void removeEdge(wEdge edge) {
		edge.node1.getOutgoingEdges().remove(edge);
		edge.node2.getIncomingEdges().remove(edge);
		allEdges.remove(edge);
	}
	
	public void removeEdges(TreeSet<wEdge> edges) {
		for (wEdge edge : edges) {
			removeEdge(edge);
		}
	}

	public TreeSet<wNode> getReachableLeafs(wNode root) {
		TreeSet<wNode> leafs = new TreeSet<wNode>();
		TreeSet<wEdge> recurse = new TreeSet<wEdge>();
		recurse.addAll(root.getOutgoingEdges());
		
		if (root.primary) leafs.add(root); // bugfix
		
		while (recurse.size() > 0) {
			wEdge check = recurse.pollFirst();
			
			if (check.node2.primary) {
				leafs.add(check.node2);
			}
			
			recurse.addAll(check.node2.getOutgoingEdges());
		}
		return leafs;
	}
	
	public TreeSet<wNode> getClosestReachablePrimaries(wNode root) {
		TreeSet<wNode> leafs = new TreeSet<wNode>();
		TreeSet<wEdge> recurse = new TreeSet<wEdge>();
				
		if (root.primary) {
			leafs.add(root);
			return leafs;
		} else {
			recurse.addAll(root.getOutgoingEdges());
		}
		
		while (recurse.size() > 0) {
			wEdge check = recurse.pollFirst();
			
			if (check.node2.primary) {
				leafs.add(check.node2);
			} else {
				recurse.addAll(check.node2.getOutgoingEdges());
			}
		}
		return leafs;
	}
	
	public HashMap<wNode, Double> getWeightedPrimaries(wNode root) {
		HashMap<wNode, Double> leafs = new HashMap<wNode, Double>();
		HashMap<wEdge, Double> recurse = new HashMap<wEdge, Double>();
				
		if (root.primary) {
			leafs.put(root, 1.0);
			return leafs;
		} else {
			for (wEdge edge : root.getOutgoingEdges()) {
				recurse.put(edge, 1.0);
			}
		}
		
		while (recurse.size() > 0) {
			wEdge check = (wEdge)recurse.keySet().toArray()[0];
			Double weight = recurse.get(check);
			recurse.remove(check);
			
			if (check.node2.primary) {
				Double w = leafs.get(check.node2);
				if (w == null) w = 0.0;
				leafs.put(check.node2, w + weight);
			} else {
				for (wEdge edge : check.node2.getOutgoingEdges()) {
					double factor = 1.0;
					
					if (check.node1.node instanceof UnaryExpression) {
						UnaryExpression e = (UnaryExpression)check.node1.node;
						if (e.op().equals(ExprOperator.CLOSURE) || e.op().equals(ExprOperator.REFLEXIVE_CLOSURE)) {
							factor = 0.2;
						}
					}
					else if (check.node1.node instanceof BinaryExpression) {
						BinaryExpression e = (BinaryExpression)check.node1.node;
						if (e.op().equals(ExprOperator.UNION) || e.op().equals(ExprOperator.INTERSECTION) || e.op().equals(ExprOperator.DIFFERENCE)) {
							factor = 0.5;
						}
						if (e.op().equals(ExprOperator.JOIN)) {
							factor = 0.7;
						}
					}
					
					Double checkVisited = recurse.get(edge);
					if (checkVisited == null) checkVisited = 0.0;
					recurse.put(edge, checkVisited + weight * factor);
				}
			}
		}
		return leafs;
	}
}
